import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from contextlib import asynccontextmanager
import router_static_html
import router_arq_upload    
from config import *
import asyncio
import threading
import os
import atexit
import time
from arq.worker import create_worker
from worker import WorkerSettings

# 全局变量存储worker任务
worker_task = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时的设置
    # 根据配置决定是否挂载静态文件目录
    if enable_static_files:
        print(f"挂载静态文件目录: {static_files_directory} -> {static_files_url}")
        app.mount(static_files_url, StaticFiles(directory=static_files_directory), name="static")
    else:
        print("静态文件目录挂载已禁用")
    
    # 启动ARQ worker
    print("启动ARQ worker...")
    asyncio.create_task(run_arq_worker())
    
    print("FastAPI应用程序准备就绪")
    yield
    
    # 关闭时的清理代码
    print("应用程序关闭中...")
    # 停止worker任务
    global worker_task
    if worker_task:
        worker_task.cancel()
        try:
            await asyncio.wait_for(worker_task, timeout=5.0)
        except (asyncio.CancelledError, asyncio.TimeoutError):
            pass
    print("清理完成")

# 在后台运行ARQ worker的函数
async def run_arq_worker():
    global worker_task
    
    try:
        print("开始启动ARQ Worker...")
        worker = create_worker(WorkerSettings)
        
        # 启动worker的任务
        worker_task = asyncio.create_task(worker.async_run())
        
        print(f"ARQ Worker 已配置，开始监听队列: {arq_queue_name}")
    except Exception as e:
        print(f"ARQ Worker 启动失败: {str(e)}")

# 创建FastAPI应用
app = FastAPI(
    title=app_name,
    description=app_description,
    version=app_version,
    lifespan=lifespan
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源，生产环境中应该设置为特定的域名
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)

app.include_router(router_arq_upload.router_upload_real, prefix=router_dict["router_upload"])
app.include_router(router_static_html.router_html_real, prefix=router_dict["router_html"])

# 添加根路径重定向到主页
@app.get("/")
async def root():
    from fastapi.responses import RedirectResponse
    # 如果静态文件目录已挂载，重定向到静态页面，否则重定向到API文档
    if enable_static_files:
        return RedirectResponse(url=f"{static_files_url}/index.html")
    else:
        return RedirectResponse(url="/docs")

# 添加API文档路径
@app.get("/api-docs")
async def api_docs():
    from fastapi.responses import RedirectResponse
    return RedirectResponse(url="/docs")

if __name__ == "__main__":
    # 启动FastAPI服务
    uvicorn.run(app, host=app_host, port=app_port) 
